package loveqq.jdbc.log.proxy;

import javassist.*;

import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * agent核心代理
 */
public class CoreProxy {

    /**
     * java.sql.PreparedStatement的实现类添加两个字段：存储SQL+参数
     */
    public static void proxyStatement(CtClass ctClass)
            throws NotFoundException, CannotCompileException {
        ClassPool pool = ctClass.getClassPool();
        // 获取全部字段，防止重复添加字段
        CtField[] allFields = ctClass.getDeclaredFields();
        Set<String> fieldNames = Stream.of(allFields)
                .map(CtField::getName)
                .collect(Collectors.toSet());
        // private String sql$agent;
        if (!fieldNames.contains(JdbcProxy.SQL_PROXY_FIELD)) {
            CtField sqlField = new CtField(pool.get("java.lang.String"), JdbcProxy.SQL_PROXY_FIELD, ctClass);
            sqlField.setModifiers(Modifier.PRIVATE);
            ctClass.addField(sqlField);
        }
        // private Map<Integer, Object> params$agent;
        if (!fieldNames.contains(JdbcProxy.PARAMS_PROXY_FIELD)) {
            CtField paramsField = new CtField(pool.get("java.util.Map"), JdbcProxy.PARAMS_PROXY_FIELD, ctClass);
            paramsField.setModifiers(Modifier.PRIVATE);
            ctClass.addField(paramsField);
        }
    }

    /**
     * java.sql.Driver的实现类复制connect(String,Properties)方法，然后将connect作为代理方法，调用复制的方法
     */
    public static void proxyDriver(CtClass ctClass) throws NotFoundException, CannotCompileException {
        // 获取方法：Connection connect(String url, java.util.Properties info)
        CtMethod connectMethod = ctClass.getMethod("connect",
                "(Ljava/lang/String;Ljava/util/Properties;)Ljava/sql/Connection;");
        if (connectMethod.isEmpty()) {
            return;
        }
        // 首先复制一个方法: connect$agent(...)
        String agentMethodName = connectMethod.getName() + "$agent";
        ctClass.addMethod(CtNewMethod.copy(connectMethod, agentMethodName, ctClass, null));
        // 用以下代码覆盖原方法
        // JdbcProxy.proxy(connect$agent(...))
        String body = String.format("{return %s.proxy(%s($$));}", JdbcProxy.class.getName(), agentMethodName);
        connectMethod.setBody(body);
    }
}
